import torch
import torch.nn as nn
import torch.nn.functional as F

from pcnets.models.builder import build_loss
class BaseClassifier(nn.Module):
    def __init__(self, loss):
        super().__init__()
        self.loss = build_loss(loss)

    def inference(self):
        raise NotImplementedError()
    
    def forward(self, data):
        """
        called by single_gpu_test, train_step
        results can be used for eval and calcuate loss
        """
        logits = self.inference(data['points'])
        data['pred_logits'] = logits
        return data

    def init_weights(self):
        pass

    def train_step(self, data, optimizer):
        """
        called by mmcv.runner
        The iteration step during training.

        This method defines an iteration step during training, except for the
        back propagation and optimizer updating, which are done in an optimizer
        hook. Note that in some complicated cases or models, the whole process
        including back propagation and optimizer updating is also defined in
        this method, such as GAN.

        Args:
            data (dict): The output of dataloader.
            optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
                runner is passed to ``train_step()``. This argument is unused
                and reserved.

        Returns:
            dict: It should contain at least 3 keys: ``loss``, ``log_vars``, \
                ``num_samples``.

                - ``loss`` is a tensor for back propagation, which can be a
                    weighted sum of multiple losses.
                - ``log_vars`` contains all the variables to be sent to the
                    logger.
                - ``num_samples`` indicates the batch size (when the model is
                    DDP, it means the batch size on each GPU), which is used for
                    averaging the logs.
        """
        gt_labels = data['label']
        logits = self.inference(data['points'])
        loss = self.loss(logits, gt_labels)

        pred_labels = logits.argmax(-1)
        acc = (gt_labels == pred_labels).sum() / pred_labels.shape[0]
        log_vars = {
            'loss': loss.item(), 
            'acc': acc.item()
        }

        outputs = dict(
            loss=loss,
            log_vars=log_vars,
            num_samples=len(data['points'])
        )

        return outputs
